package io.kwdb.sdk;

import static org.junit.jupiter.api.Assertions.*;

import io.kwdb.sdk.internal.Batcher;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

/**
 * Test cases for the Batcher component. Validates batching logic, triggers, and backpressure
 * handling.
 */
public class BatcherTest {

  private Batcher<String> batcher;
  private List<List<String>> submittedBatches;
  private CountDownLatch batchLatch;

  @BeforeEach
  public void setup() {
    submittedBatches = Collections.synchronizedList(new ArrayList<>());
    batchLatch = new CountDownLatch(1);
  }

  @AfterEach
  public void cleanup() throws Exception {
    if (batcher != null) {
      batcher.close();
    }
  }

  @Test
  public void testBatchSizeTrigger() throws Exception {
    // Create batcher with max batch size of 5
    batcher =
        new Batcher<>(
            5,
            10000,
            100,
            Config.Backpressure.BLOCK,
            batch -> {
              submittedBatches.add(new ArrayList<>(batch));
              batchLatch.countDown();
            });

    // Add 5 items - should trigger batch
    for (int i = 0; i < 5; i++) {
      assertTrue(batcher.offer("item" + i, 1000));
    }

    // Wait for batch submission
    assertTrue(batchLatch.await(2, TimeUnit.SECONDS));

    // Verify batch
    assertEquals(1, submittedBatches.size());
    assertEquals(5, submittedBatches.get(0).size());
    assertEquals("item0", submittedBatches.get(0).get(0));
    assertEquals("item4", submittedBatches.get(0).get(4));
  }

  @Test
  public void testLingerTimeTrigger() throws Exception {
    // Create batcher with large batch size but short linger time
    batcher =
        new Batcher<>(
            100,
            100,
            1000,
            Config.Backpressure.BLOCK,
            batch -> {
              submittedBatches.add(new ArrayList<>(batch));
              batchLatch.countDown();
            });

    // Add only 3 items (won't hit batch size)
    for (int i = 0; i < 3; i++) {
      assertTrue(batcher.offer("item" + i, 1000));
    }

    // Wait for linger time to trigger
    assertTrue(batchLatch.await(500, TimeUnit.MILLISECONDS));

    // Verify batch
    assertEquals(1, submittedBatches.size());
    assertEquals(3, submittedBatches.get(0).size());
  }

  @Test
  public void testFlush() throws Exception {
    batcher =
        new Batcher<>(
            100,
            10000,
            1000,
            Config.Backpressure.BLOCK,
            batch -> {
              submittedBatches.add(new ArrayList<>(batch));
              batchLatch.countDown();
            });

    // Add items that won't trigger batch
    for (int i = 0; i < 10; i++) {
      assertTrue(batcher.offer("item" + i, 1000));
    }

    // Flush should trigger immediate submission
    batcher.flush();

    // Wait for batch
    assertTrue(batchLatch.await(1, TimeUnit.SECONDS));

    // Verify batch
    assertEquals(1, submittedBatches.size());
    assertEquals(10, submittedBatches.get(0).size());
  }

  @Test
  public void testBlockBackpressure() throws Exception {
    AtomicBoolean processingBlocked = new AtomicBoolean(true);
    CountDownLatch processingLatch = new CountDownLatch(1);

    // Create batcher with small queue
    batcher =
        new Batcher<>(
            2,
            10000,
            4,
            Config.Backpressure.BLOCK,
            batch -> {
              // Block processing until flag is cleared
              while (processingBlocked.get()) {
                try {
                  Thread.sleep(10);
                } catch (InterruptedException e) {
                  break;
                }
              }
              submittedBatches.add(new ArrayList<>(batch));
              processingLatch.countDown();
            });

    // Fill the queue
    assertTrue(batcher.offer("item1", 1000));
    assertTrue(batcher.offer("item2", 1000)); // Triggers batch, but processing is blocked
    assertTrue(batcher.offer("item3", 1000));
    assertTrue(batcher.offer("item4", 1000)); // Queue is now full

    // Next offer should block and timeout
    long start = System.currentTimeMillis();
    assertFalse(batcher.offer("item5", 100)); // Should timeout
    long elapsed = System.currentTimeMillis() - start;
    assertTrue(elapsed >= 90, "Should have blocked for ~100ms");

    // Unblock processing
    processingBlocked.set(false);
    assertTrue(processingLatch.await(1, TimeUnit.SECONDS));

    // Now we should be able to add more
    assertTrue(batcher.offer("item5", 1000));
  }

  @Test
  public void testDropOldestBackpressure() throws Exception {
    AtomicInteger dropCount = new AtomicInteger(0);

    // Create batcher with DROP_OLDEST and small queue
    batcher =
        new Batcher<>(
            2,
            10000,
            4,
            Config.Backpressure.DROP_OLDEST,
            batch -> {
              // Slow processing to cause backpressure
              try {
                Thread.sleep(100);
              } catch (InterruptedException e) {
                // Ignore
              }
              submittedBatches.add(new ArrayList<>(batch));
            });

    // Rapidly add many items
    for (int i = 0; i < 20; i++) {
      boolean accepted = batcher.offer("item" + i, 100);
      if (!accepted) {
        dropCount.incrementAndGet();
      }
    }

    // Some items should have been dropped
    assertTrue(dropCount.get() > 0, "Should have dropped some items");

    // Flush remaining
    batcher.flush();
    Thread.sleep(200); // Wait for processing

    // Total processed should be less than 20
    int totalProcessed = submittedBatches.stream().mapToInt(List::size).sum();
    assertTrue(totalProcessed < 20, "Some items should have been dropped");
  }

  @Test
  public void testMultipleProducers() throws Exception {
    CountDownLatch startLatch = new CountDownLatch(1);
    AtomicInteger successCount = new AtomicInteger(0);

    batcher =
        new Batcher<>(
            10,
            50,
            100,
            Config.Backpressure.BLOCK,
            batch -> {
              submittedBatches.add(new ArrayList<>(batch));
            });

    // Create multiple producer threads
    int numThreads = 5;
    int itemsPerThread = 20;
    ExecutorService executor = Executors.newFixedThreadPool(numThreads);

    for (int t = 0; t < numThreads; t++) {
      final int threadId = t;
      executor.submit(
          () -> {
            try {
              startLatch.await();
              for (int i = 0; i < itemsPerThread; i++) {
                if (batcher.offer("thread" + threadId + "-item" + i, 1000)) {
                  successCount.incrementAndGet();
                }
              }
            } catch (Exception e) {
              e.printStackTrace();
            }
          });
    }

    // Start all threads
    startLatch.countDown();

    // Wait for completion
    executor.shutdown();
    assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));

    // Flush remaining
    batcher.flush();
    Thread.sleep(100);

    // Verify all items were processed
    int totalProcessed = submittedBatches.stream().mapToInt(List::size).sum();
    assertEquals(successCount.get(), totalProcessed);
    assertEquals(numThreads * itemsPerThread, totalProcessed);
  }

  @Test
  public void testCloseFlushesRemaining() throws Exception {
    batcher =
        new Batcher<>(
            100,
            10000,
            1000,
            Config.Backpressure.BLOCK,
            batch -> {
              submittedBatches.add(new ArrayList<>(batch));
            });

    // Add items that won't trigger batch
    for (int i = 0; i < 7; i++) {
      assertTrue(batcher.offer("item" + i, 1000));
    }

    // Close should flush remaining items
    batcher.close();

    // Verify items were flushed
    assertEquals(1, submittedBatches.size());
    assertEquals(7, submittedBatches.get(0).size());

    // Should not accept new items after close
    assertFalse(batcher.offer("afterClose", 100));
  }

  @Test
  public void testEmptyBatchNotSubmitted() throws Exception {
    AtomicInteger submitCount = new AtomicInteger(0);

    batcher =
        new Batcher<>(
            10,
            100,
            100,
            Config.Backpressure.BLOCK,
            batch -> {
              submitCount.incrementAndGet();
              assertFalse(batch.isEmpty(), "Empty batch should not be submitted");
            });

    // Just wait for linger time without adding items
    Thread.sleep(200);

    // No batch should have been submitted
    assertEquals(0, submitCount.get());

    // Flush with no items should also not submit
    batcher.flush();
    Thread.sleep(50);
    assertEquals(0, submitCount.get());
  }

  @Test
  public void testBatcherMetrics() throws Exception {
    io.kwdb.sdk.internal.Metrics metrics = io.kwdb.sdk.internal.Metrics.getInstance();
    metrics.reset();

    batcher =
        new Batcher<>(
            5,
            1000,
            100,
            Config.Backpressure.DROP_OLDEST,
            batch -> {
              // Simulate processing
              try {
                Thread.sleep(10);
              } catch (InterruptedException e) {
                // Ignore
              }
            });

    // Add items to trigger batch
    for (int i = 0; i < 10; i++) {
      batcher.offer("item" + i, 1000);
    }

    // Wait for processing
    Thread.sleep(100);

    // Check metrics were updated
    io.kwdb.sdk.internal.Metrics.MetricsSnapshot snapshot = metrics.getSnapshot();
    assertTrue(snapshot.currentBatchSize > 0, "Should have recorded batch size");
  }
}
